# 支持多张图 合成
model_name = '/data/models/Qwen-Image-Edit-2509/'

from util_for_argparse import get_cuda_port_args

args = get_cuda_port_args()

'''prompt
换成白色上衣短袖，白色下身短裤
'''

import os
import torch
from PIL import Image
from diffusers import QwenImageEditPlusPipeline

pipeline = QwenImageEditPlusPipeline.from_pretrained(model_name, device_map='balanced', torch_type=torch.bfloat16)
print("pipeline loaded")

# pipeline.to('cuda')
pipeline.set_progress_bar_config(disable=None)
image1 = Image.open("tmp_tryon.jpg")
image2 = Image.open("tmp_cloth.jpg")
prompt = "换成白色上衣短袖，白色下身短裤"
inputs = {
    "image": [image1],
    "prompt": prompt,
    "generator": torch.manual_seed(0),
    "true_cfg_scale": 4.0,
    "negative_prompt": " ",
    "num_inference_steps": 40,
    "guidance_scale": 1.0,
    "num_images_per_prompt": 1,
}
with torch.inference_mode():
    output = pipeline(**inputs)
    output_image = output.images[0]
    output_image.save("tmp_output_image_edit_plus.jpg")
    print("image saved at", os.path.abspath("tmp_output_image_edit_plus.jpg"))
